import torch
import torch.nn as nn
import torch.nn.functional as F

import torchvision
import torchvision.datasets as dset
import torchvision.transforms as transforms

import torch.optim as optim
import torchvision.models as models

import PIL.Image as Image

image_size = (224, 224)
data_transform = transforms.Compose([transforms.Resize((224, 224)),
                                     #transforms.RandomHorizontalFlip(),
                                     # T.RandomVerticalFlip(),
                                     # T.ColorJitter(0.5, 0.5, 0.5, 0.5),
                                     #transforms.Pad(10),
                                     #transforms.RandomCrop((224, 224)),
                                     # T.RandomRotation(90),
                                     transforms.ToTensor(),
                                     transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])

train_data = dset.ImageFolder(root="data", transform=data_transform)
# 数据集长度
totallen = len(train_data)
print('train data length:', totallen)

trainlen = int(totallen * 0.7)
vallen = totallen - trainlen
train_db, val_db = torch.utils.data.random_split(train_data, [trainlen, vallen])
print('train:', len(train_db), 'validation:', len(val_db))

# batch size
bs = 8
# 训练集
train_loader = torch.utils.data.DataLoader(train_db, batch_size=bs, shuffle=True)
# 验证集
val_loader = torch.utils.data.DataLoader(val_db, batch_size=bs, shuffle=True)


def get_num_correct(out, labels):
    return out.argmax(dim=1).eq(labels).sum().item()


resnet50 = models.resnet50(pretrained=True)
model = resnet50
n_classes = len(train_data.classes)
model.fc = nn.Linear(2048, n_classes)


import torch.nn.init as init

for name, module in model._modules.items():
    if (name == 'fc'):
        # print(module.weight.shape)
        init.kaiming_uniform_(module.weight, a=0, mode='fan_in')

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(device)


optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
epoch_num = 9
model = model.to(device)
if hasattr(torch.cuda, 'empty_cache'):
    torch.cuda.empty_cache()

if __name__ == '__main__':
    for epoch in range(epoch_num):
        total_loss = 0
        total_correct = 0
        val_correct = 0

        if hasattr(torch.cuda, 'empty_cache'):
            torch.cuda.empty_cache()

        for batch in train_loader:  # GetBatch
            images, labels = batch
            images = images.to(device)
            labels = labels.to(device)
            #print(images.size())
            outs = model(images)  # PassBatch
            loss = F.cross_entropy(outs, labels)  # CalculateLoss
            optimizer.zero_grad()
            loss.backward()  # CalculateGradients
            optimizer.step()  # UpdateWeights
            total_loss += loss.item()
            total_correct += get_num_correct(outs, labels)

            if hasattr(torch.cuda, 'empty_cache'):
                torch.cuda.empty_cache()
        with torch.no_grad():
            for batch in val_loader:
                images, labels = batch
                images = images.to(device)
                labels = labels.to(device)
                outs = model(images)
                val_correct += get_num_correct(outs, labels)
        print("loss:", total_loss, "train_correct:", total_correct / trainlen, "val_correct:", val_correct / vallen)



    torch.save(model, 'falldetect.pth')